"""
ConnLLM Moonshot部分输出能力V2实现
提供Moonshot API部分输出功能
"""

from typing import Dict, Any, List, Optional
import logging
import copy

from ...core_v2.capability import ModelCapability

logger = logging.getLogger("connllm.implementations.moonshot.partial_mode_v2")

class MoonshotPartialModeV2(ModelCapability):
	"""Moonshot部分输出能力V2实现"""
	
	def __init__(self, provider):
		"""
		初始化能力
		
		Args:
			provider: 提供商实例
		"""
		super().__init__(provider)
		self.current_messages = None  # 当前上下文
		self.last_message_index = -1  # 最后一条消息索引
	
	def continue_from_partial(self, text: str, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
		"""
		从部分输出继续生成
		
		Args:
			text: 已生成的部分输出
			messages: 消息列表
			
		Returns:
			API响应
		"""
		# 保存当前消息上下文
		self.current_messages = copy.deepcopy(messages)
		self.last_message_index = len(messages) - 1
		
		# 检查是否需要新建上下文
		if not self._is_same_context(messages):
			logger.debug("上下文已变更，创建新的部分输出上下文")
			self.current_messages = copy.deepcopy(messages)
			self.last_message_index = len(messages) - 1
		
		# 准备继续生成的参数
		params = {
			"model": self.provider.config_manager.get_model(),
			"messages": self.current_messages,
			"suffix": "",
			"stream": False
		}
		
		# 添加继续生成参数
		completion_settings = self.provider.config_manager.get_completion_settings()
		params.update({
			"max_tokens": completion_settings.get("max_tokens", 4096),
			"temperature": completion_settings.get("temperature", 0.7),
			"top_p": completion_settings.get("top_p", 1.0),
			"presence_penalty": completion_settings.get("presence_penalty", 0.0),
			"frequency_penalty": completion_settings.get("frequency_penalty", 0.0),
		})
		
		# 如果当前上下文最后一条消息是助手消息，且内容与部分输出匹配
		if (self.last_message_index >= 0 and 
			self.current_messages[self.last_message_index].get("role") == "assistant" and
			self.current_messages[self.last_message_index].get("content", "").startswith(text)):
			# 直接使用已有上下文，不做修改
			logger.debug("继续已有的部分输出")
		else:
			# 添加部分输出作为助手消息
			if self.last_message_index >= 0 and self.current_messages[self.last_message_index].get("role") == "assistant":
				# 更新最后一条助手消息
				self.current_messages[self.last_message_index]["content"] = text
			else:
				# 添加新的助手消息
				self.current_messages.append({
					"role": "assistant",
					"content": text
				})
				self.last_message_index = len(self.current_messages) - 1
			
			logger.debug(f"添加部分输出为助手消息: {text[:50]}...")
		
		# 发送API请求
		endpoint = self.provider._get_completion_endpoint()
		
		try:
			# 创建请求上下文
			context = {"original_messages": messages, "partial_mode": True}
			
			# 应用请求处理器
			for processor in self.provider.request_processors:
				params = processor.process_request(params, context)
			
			# 发送HTTP请求
			response = self.provider.http_client.request(
				method="POST",
				endpoint=endpoint,
				data=params
			)
			result = response.json()
			
			# 应用响应处理器
			for processor in self.provider.response_processors:
				result = processor.process_response(result, context)
			
			# 更新当前上下文，添加响应作为助手消息
			if "choices" in result and len(result["choices"]) > 0:
				choice = result["choices"][0]
				
				if "message" in choice and "content" in choice["message"]:
					content = choice["message"]["content"]
					
					if self.last_message_index >= 0 and self.current_messages[self.last_message_index].get("role") == "assistant":
						# 更新最后一条助手消息
						self.current_messages[self.last_message_index]["content"] = content
					else:
						# 添加新的助手消息
						self.current_messages.append({
							"role": "assistant",
							"content": content
						})
						self.last_message_index = len(self.current_messages) - 1
				
			return result
			
		except Exception as e:
			logger.error(f"部分输出继续生成失败: {str(e)}")
			raise e
	
	def get_current_messages(self) -> List[Dict[str, Any]]:
		"""
		获取当前消息上下文
		
		Returns:
			当前消息列表
		"""
		return self.current_messages if self.current_messages else []
	
	def _is_same_context(self, messages: List[Dict[str, Any]]) -> bool:
		"""
		检查是否是同一个上下文
		
		Args:
			messages: 消息列表
			
		Returns:
			是否是同一个上下文
		"""
		if not self.current_messages:
			return False
		
		# 简单检查：长度和用户最后一条消息内容
		if len(messages) != len(self.current_messages):
			return False
		
		# 找到最后一条用户消息
		last_user_msg = None
		last_current_user_msg = None
		
		for i in range(len(messages) - 1, -1, -1):
			if i < len(messages) and messages[i].get("role") == "user":
				last_user_msg = messages[i].get("content", "")
				break
				
		for i in range(len(self.current_messages) - 1, -1, -1):
			if i < len(self.current_messages) and self.current_messages[i].get("role") == "user":
				last_current_user_msg = self.current_messages[i].get("content", "")
				break
		
		return last_user_msg == last_current_user_msg
